// Copyright 2011 The Apache Software Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package org.apache.tapestry5.csrfprotection;

import java.lang.reflect.Method;
import java.util.Collection;
import java.util.Comparator;
import java.util.List;

import org.apache.tapestry5.ComponentResources;
import org.apache.tapestry5.Link;
import org.apache.tapestry5.MarkupWriter;
import org.apache.tapestry5.corelib.base.AbstractLink;
import org.apache.tapestry5.corelib.components.AjaxFormLoop;
import org.apache.tapestry5.corelib.components.BeanEditForm;
import org.apache.tapestry5.corelib.components.Form;
import org.apache.tapestry5.corelib.components.FormInjector;
import org.apache.tapestry5.corelib.components.RemoveRowLink;
import org.apache.tapestry5.corelib.data.InsertPosition;
import org.apache.tapestry5.corelib.internal.AjaxFormLoopContext;
import org.apache.tapestry5.csrfprotection.CsrfProtected;
import org.apache.tapestry5.csrfprotection.CsrfException;
import org.apache.tapestry5.csrfprotection.CsrfProtectable;
import org.apache.tapestry5.csrfprotection.services.CsrfProtectedPages;
import org.apache.tapestry5.csrfprotection.services.CsrfProtectionModule;
import org.apache.tapestry5.dom.Element;
import org.apache.tapestry5.internal.services.RequestConstants;
import org.apache.tapestry5.ioc.annotations.Value;
import org.apache.tapestry5.json.JSONObject;
import org.apache.tapestry5.model.MutableComponentModel;
import org.apache.tapestry5.plastic.FieldHandle;
import org.apache.tapestry5.plastic.MethodAdvice;
import org.apache.tapestry5.plastic.MethodDescription;
import org.apache.tapestry5.plastic.MethodInvocation;
import org.apache.tapestry5.plastic.PlasticClass;
import org.apache.tapestry5.plastic.PlasticField;
import org.apache.tapestry5.plastic.PlasticMethod;
import org.apache.tapestry5.services.ApplicationStateManager;
import org.apache.tapestry5.services.ClientBehaviorSupport;
import org.apache.tapestry5.services.ComponentClassResolver;
import org.apache.tapestry5.services.FormSupport;
import org.apache.tapestry5.services.Request;
import org.apache.tapestry5.services.javascript.JavaScriptSupport;
import org.apache.tapestry5.services.transform.ComponentClassTransformWorker2;
import org.apache.tapestry5.services.transform.TransformationSupport;
import org.jaxen.JaxenException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.formos.tapestry.xpath.TapestryXPath;

/**
 * This class introduces code for cross-site request forgery protection at runtime. 
 * Components and pages of the Tapestry project that uses this module are scanned and modified. The core components of Tapestry are 
 * also in scope of the class transformation. 
 * 
 * The cross-site request forgery protection works with two modes which can be configured in the Tapestry module class:
 * <ul>
 * 		<li>Explicit mode</li>
 * 		<li>Auto-protection mode</li>
 * </ul>
 * 
 * <h2>Explicit mode</h2>
 * <p>
 * In the explicit mode components can be protected by using the instance based mixin {@link org.apache.tapestry5.csrfprotection.mixins.CsrfProtected} or on a page level
 * by using the {@link CsrfProtected} annotation. The instance based mixin calls the insertCSRFToken method for the component. This method is added
 * to all components that are available in controlled packages. There are different types of this method and one can override it simply by providing this 
 * method in a component, but there is a generic version that should work for all components that rely on HTML rendered form and anchor tags. Event handling 
 * methods should be marked with the {@link CsrfProtected} annotation. A cross-site request forgery token check is then added as advice into the event handling 
 * method. If the event handling method is not present, because an embedded component is used then the {@link CsrfProtected} annotation should be placed on the
 * page. In this case all events that occur on that page are protected. All pages that are protected are stored in the {@link CsrfProtectedPages} service and 
 * evaluated in the {@link CsrfProtectionFilter}.
 * </p>
 * <h2>Auto mode</h2>
 * <p>In the auto mode a decorator for the ComponentEventLinkEncoder is used, which adds the cross-site request forgery protection token to all components. The
 * {@link CsrfProtectionFilter} checks the component events then for all requests. In this case this transformation worker looks for pages that are marked 
 * with the ${@link NotCsrfPortected}. These pages are also stored in the {@link CsrfProtectedPages} service and evaluated by the {@link CsrfProtectionFilter}. 
 * A token check is in that case not executed.   
 * 
 * @see CsrfProtectionFilter
 * @see CsrfProtectedPages
 */
public class CsrfProtectionClassTransformWorker implements
		ComponentClassTransformWorker2 {
	private Request request;

	private final ApplicationStateManager applicationStateManager;
	private final ComponentClassResolver componentClassResolver;
	private final CsrfProtectedPages csrfProtectedPages;

	private String antiCsrfMode;
	
	private Logger logger = LoggerFactory
			.getLogger(CsrfProtectionClassTransformWorker.class);

	/**
	 * Creates a new instance and injects the required services and configuration parameters.
	 * @param request
	 * @param applicationStateManager
	 * @param logger
	 * @param componentClassResolver
	 * @param antiCsrfMode
	 * @param antiCsrfTokenType
	 * @param antiCsrfTokenPersistence
	 * @param csrfProtectedPages
	 */
	public CsrfProtectionClassTransformWorker(
			Request request,
			ApplicationStateManager applicationStateManager,
			Logger logger,
			ComponentClassResolver componentClassResolver,
			@Value("${" + CsrfProtectionModule.ANTI_CSRF_MODE + "}") String antiCsrfMode,			
			CsrfProtectedPages csrfProtectedPages) {
		this.request = request;
		this.applicationStateManager = applicationStateManager;
		this.componentClassResolver = componentClassResolver;
		this.antiCsrfMode = antiCsrfMode;
		this.csrfProtectedPages = csrfProtectedPages;
	}

	/**
	 * Performs the cross-site request forgery related class transformation.
	 * @param plasticClass Current class that is going to be transformed
	 * @param support
	 * @param model
	 */
	public void transform(PlasticClass plasticClass,
			TransformationSupport support, MutableComponentModel model) {
		logger.debug("CsrfProtectionWorker mode: " + antiCsrfMode);
		/*
		 * Create the insertCSRFToken method for all components in all
		 * controlled package.
		 */
		if (!CsrfProtectionModule.ANTI_CSRF_MODE_OFF.equals(antiCsrfMode)) {
			logger.debug("Csrf protection check class: "
					+ plasticClass.getClassName());

			instrumentTokenCheck(plasticClass);

			try {
				instrumentTokenInject(plasticClass);
			} catch (ClassNotFoundException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
	}

	/**
	 * Method advice that is added to CsrfProtected annotated methods.
	 */
	private final MethodAdvice csrfTokenCheckAdvice = new MethodAdvice() {
		public void advise(MethodInvocation invocation) {			
			try {
				CsrfTokenProvider.checkToken(request, applicationStateManager);
				invocation.proceed();
			} catch (CsrfException e) {
				invocation.setCheckedException(e);
				invocation.rethrow();
			}
		}
	};

	/**
	 * This methods protects classes and methods that are annotated with the
	 * Protected annotation. The according Tapestry elements are Page classes
	 * and event handler method for which these protection is intended. The HTTP
	 * request parameter is compared against the session attribute for the CSRF
	 * token. If the values are not matching or not present an exception is
	 * thrown. The annotation should be used on Tapestry Page classes and on
	 * event handling methods.
	 * 
	 * @param plasticClass
	 */
	private void instrumentTokenCheck(final PlasticClass plasticClass) {
		logger.debug("instrumentTokenCheck checks: "
				+ plasticClass.getClassName());

		protectEventHandlers(plasticClass, csrfTokenCheckAdvice);
		protectPage(plasticClass, csrfTokenCheckAdvice);
	}

	/**
	 * Adds the token check advice to the page that has the @CsrfProtected annoation.
	 * @param plasticClass
	 * @param csrfTokenCheckAdvice
	 */
	private void protectPage(PlasticClass plasticClass,
			MethodAdvice csrfTokenCheckAdvice) {
		logger.debug("protectPage checks " + plasticClass.getClassName());

		if (plasticClass.hasAnnotation(CsrfProtected.class)) {
			// remember page, for the CsrfProtectionFilter
			logger.debug("recording protected page: "
					+ plasticClass.getClassName());
			String pageName = componentClassResolver
					.resolvePageClassNameToPageName(plasticClass.getClassName());
			csrfProtectedPages.addPageName(pageName);
		}

		if (plasticClass.hasAnnotation(NotCsrfProtected.class)) {
			// remember page, for the CsrfProtectionFilter
			logger.debug("recording unprotected page: "
					+ plasticClass.getClassName());
			String pageName = componentClassResolver
					.resolvePageClassNameToPageName(plasticClass.getClassName());
			csrfProtectedPages.addUnprotectedPageName(pageName);
		}
	}

	
	/**
	 * Protects all event handlers that are marked with the @CsrfProtected annotation.
	 * @param plasticClass
	 * @param csrfTokenCheckAdvice
	 */
	private void protectEventHandlers(PlasticClass plasticClass,
			MethodAdvice csrfTokenCheckAdvice) {
		// Method adivce for event handlers

		for (final PlasticMethod method : plasticClass
				.getMethodsWithAnnotation(CsrfProtected.class)) {
			method.addAdvice(csrfTokenCheckAdvice);
		}
	}

	/**
	 * For the different component types the according insertCSRFtoken methods
	 * are added.
	 * @param plasticClass
	 * @throws ClassNotFoundException
	 */
	private void instrumentTokenInject(PlasticClass plasticClass)
			throws ClassNotFoundException {
		/*
		 * Create the method description of insertCSRFToken and the advises
		 * according to the component types. If there is no special way to
		 * insert a CSRF token a generic method is used.
		 */

		/* Generic insert CSRF token method */
		MethodDescription insertCSRFToken = new MethodDescription("void",
				"insertCSRFToken", MarkupWriter.class.getName(),
				String.class.getName());

		// Check the class name and decide which method to inject that should
		// insert the CSRF token
		String className = plasticClass.getClassName();

		boolean inControlledPackage = isInControlledPackage(className);
		if (inControlledPackage) {

			// Check if component has already a insertCSRFToken method
			if (!plasticClass.isMethodImplemented(insertCSRFToken)) {
				PlasticMethod method = plasticClass
						.introduceMethod(insertCSRFToken);

				if (Class.forName(plasticClass.getClassName()).equals(
						Form.class)) {
					method.addAdvice(getFormAdvice(plasticClass));
				} else 	if (Class.forName(plasticClass.getClassName()).equals(
						AbstractLink.class)) {
					method.addAdvice(getLinkAdvice(plasticClass));
				} else if (Class.forName(plasticClass.getClassName()).equals(
						BeanEditForm.class)) {
					method.addAdvice(getGenericAdvice());
				} 
				else if (Class.forName(plasticClass.getClassName()).equals(
						RemoveRowLink.class)){
					method.addAdvice(getRemoveRowAdivce(plasticClass));
				} 
				else if (Class.forName(plasticClass.getClassName()).equals(
						FormInjector.class)){
					method.addAdvice(getFormInjectorAdvice(plasticClass));
				}
				else if (Class.forName(plasticClass.getClassName()).equals(
						AjaxFormLoop.class)){ 
					method.addAdvice(getAjaxFormLoopAdvice(plasticClass));
				}
				else {
					method.addAdvice(getGenericAdvice());
				}
			}

			plasticClass.introduceInterface(CsrfProtectable.class);
		}
	}

	/**
	 * This advice calls the insertCSRFToken method provided by the FormInjector (also added through 
	 * class transformation). This protection only takes care about the AddRowLink created javascript 
	 * event links. The RemoveRowLink is protected separately.
	 * @param plasticClass
	 * @return MethodAdvice 
	 */
	private MethodAdvice getAjaxFormLoopAdvice(PlasticClass plasticClass) {
		final FieldHandle _rowInjector = getFieldHandleByName(plasticClass, "rowInjector");
		
		return new MethodAdvice(){			 
			public void advise(MethodInvocation invocation) {			
				String token = (String) invocation.getParameter(1);
				Object rowInjector = _rowInjector.get(invocation.getInstance());
				
				if(rowInjector.getClass().getName().equals(FormInjector.class.getName())){			  
			       
			       	if(rowInjector instanceof CsrfProtectable){
			       		
			       		((CsrfProtectable) rowInjector).insertCSRFToken((MarkupWriter)invocation.getParameter(0), token);
			       	}
			    }						
			}
		};	      
	}

	/**
	 * This advice is required to protect the AddRowLink component which is another component that is embedded
	 * in the AjaxFormLoop component.
	 * @param plasticClass
	 * @return MethodAdvice 
	 */
	private MethodAdvice getFormInjectorAdvice(PlasticClass plasticClass) {
		final FieldHandle _resources = getFieldHandleByName(plasticClass,"resources");
		final FieldHandle _clientBehaviorSupport = getFieldHandleByName(plasticClass,"clientBehaviorSupport");
		final FieldHandle _context = getFieldHandleByName(plasticClass,"context");
		final FieldHandle _formSupport = getFieldHandleByName(plasticClass, "formSupport");
		final FieldHandle _clientId = getFieldHandleByName(plasticClass, "clientId");
		final FieldHandle _position = getFieldHandleByName(plasticClass, "position");
		final FieldHandle _show = getFieldHandleByName(plasticClass, "show");
		
		return new MethodAdvice() {
			public void advise(MethodInvocation invocation) {

				/*
				 * Compatibility to Tapestry 5.3.0
				 */

				boolean oldTapestry = true;
				Method[] methods = ClientBehaviorSupport.class.getMethods();
				for (Method method : methods) {
					if (method.getName().equals("updateFormInjector")) {
						oldTapestry = false;
					}
				}
				
				if(oldTapestry){
					logger.warn("CSRF protection module cannot protect the AjaxFormLoop component in explicit mode with this version of Tapestry.");
					return;
				}

				ComponentResources resources = (ComponentResources) _resources
						.get(invocation.getInstance());
				ClientBehaviorSupport clientBehaviorSupport = (ClientBehaviorSupport) _clientBehaviorSupport
						.get(invocation.getInstance());
				Object[] context = (Object[]) _context.get(invocation
						.getInstance());
				String token = (String) invocation.getParameter(1);
				FormSupport formSupport = (FormSupport) _formSupport
						.get(invocation.getInstance());
				String clientId = (String) _clientId.get(invocation
						.getInstance());
				InsertPosition position = (InsertPosition) _position
						.get(invocation.getInstance());
				String show = (String) _show.get(invocation.getInstance());

				Link link = resources.createEventLink(
						FormInjector.INJECT_EVENT, context);
				link.addParameter(RequestConstants.FORM_CLIENTID_PARAMETER,
						formSupport.getClientId());
				link.addParameter(RequestConstants.FORM_COMPONENTID_PARAMETER,
						formSupport.getFormComponentId());
				link.addParameter(CsrfTokenProvider.TOKEN_NAME, token);

				clientBehaviorSupport.updateFormInjector(clientId, link,
						position, show);
			}
		};
	}
		
	/**
	 * Special insertCSRFToken advice for the RemoveRow component. (Required to protect the AjaxFormLoop component).
	 * @param plasticClass
	 * @return
	 */
	private MethodAdvice getRemoveRowAdivce(PlasticClass plasticClass) {
		final FieldHandle _jsSupport = getFieldHandleByName(plasticClass, "jsSupport");
		final FieldHandle _clientId  = getFieldHandleByName(plasticClass, "clientId");
		final FieldHandle _context = getFieldHandleByName(plasticClass, "context");
		return new MethodAdvice(){
			public void advise(MethodInvocation invocation){				
				/*
				 * If these fields are not available the CSRF protection is used for a previous version
				 * of Tapestry.
				 */
				if(_clientId == null){
					return;
				}
				
				String clientId = (String) _clientId.get(invocation.getInstance());
				AjaxFormLoopContext context = (AjaxFormLoopContext) _context.get(invocation.getInstance());				  
				JavaScriptSupport jsSupport = (JavaScriptSupport)_jsSupport.get(invocation.getInstance());
										
				// get token
				String token = (String) invocation.getParameter(1);
				Link link = context.getResources().createEventLink("triggerRemoveRow", context.toClientValue());
		        link.addParameter("t:antiCsrfToken", token);

		        String asURI = link.toURI();
		        		  
		        JSONObject spec = new JSONObject();
		        spec.put("link", clientId);
		        spec.put("fragment", context.currentFragmentId());
		        spec.put("url", asURI);

		        jsSupport.updateInitializerCall("formLoopRemoveLink", spec, new Comparator<JSONObject>() {
		            public int compare(JSONObject o1, JSONObject o2) {
		                if (o1.has("link") && o1.has("fragment") && o2.has("link") && o2.has("fragment")) {
		                    if (o1.get("link").equals(o2.get("link")) && o1.get("fragment").equals(o2.get("fragment"))) {
		                        return 0;
		                    }
		                }
		                return 1;
		            }

		            @Override
		            public boolean equals(Object obj) {
		                return false;
		            }
		        });
			}
		};
	}

	/**
	 * Creates a MethodAdivce for the Tapestry Form component, which injects the
	 * CSRF protection token. This is tight coupled to the implementation of
	 * Form class and requires an instance variable form of type Element.
	 * 
	 * @param plasticClass
	 * @return
	 */
	private MethodAdvice getFormAdvice(PlasticClass plasticClass) {
		return createFormAdvice(getFieldHandleByName(plasticClass, "form"), getFieldHandleByName(plasticClass, "javascriptSupport"));
		
	}

	/**
	 * This method is a helper for getFormAdvice and does the byte code
	 * modification with Plastic.
	 * 
	 * @param formHandle
	 * @return
	 */
	private MethodAdvice createFormAdvice(final FieldHandle formHandle, final FieldHandle javascriptSupportHandle) {
		
		return new MethodAdvice() {
			public void advise(MethodInvocation invocation) {
				// Form form = (Form) invocation.getInstance();
				Element formElement = (Element) formHandle.get(invocation
						.getInstance());				
					
				// get token
				String token = (String) invocation.getParameter(1);

				// insert hidden form element with token as value
				addHiddenTokenToForm(formElement, token);						
			}
		};
	}

	/**
	 * Advice to add a hidden cross-site request forgery protection token to a form.
	 * @param formElement
	 * @param token
	 */
	private void addHiddenTokenToForm(Element formElement, String token) {
		formElement.element("input", "type", "hidden", "name",
				CsrfTokenProvider.TOKEN_NAME, "value", token);
	}

	/**
	 * Advice to add the token as parameter to a link.
	 * @param plasticClass
	 * @return
	 */
	private MethodAdvice getLinkAdvice(PlasticClass plasticClass) {
		return createLinkAdvice(getFieldHandleByName(plasticClass, "element"));

	}

	/**
	 * Creates the link advice
	 * @param elementHandle
	 * @return
	 */
	private MethodAdvice createLinkAdvice(FieldHandle elementHandle) {
		final FieldHandle _elementHandle = elementHandle;
		return new MethodAdvice() {
			public void advise(MethodInvocation invocation) {
				// TODO Auto-generated method stub
				Element element = (Element) _elementHandle.get(invocation
						.getInstance());
				// get token
				String token = (String) invocation.getParameter(1);
				if (element != null)
					addParameterTokenToLink(element, token);
			}
		};
	}

	
	/**
	 * Adds the parameter to the link.
	 * @param linkElement
	 * @param token
	 */
	private void addParameterTokenToLink(Element linkElement, String token) {
		// insert link parameter		
		if (linkElement.getAttribute("href") != null) {
			StringBuilder href = new StringBuilder(
					linkElement.getAttribute("href"));
			if (href.indexOf("?") > 0) { // there are already some parameters
				href = href.append("&" + CsrfTokenProvider.TOKEN_NAME + "=")
						.append(token);
			} else { // add parameter
				href = href.append("?" + CsrfTokenProvider.TOKEN_NAME + "=")
						.append(token);
			}
			linkElement.forceAttributes("href", href.toString());
		}
	}

	/**
	 * The generic method advice to insert a cross-site request forgery protection token uses the Tapestry XPath for that tasks.
	 * @return
	 */
	private MethodAdvice getGenericAdvice() {
		return new MethodAdvice() {
			public void advise(MethodInvocation invocation) {
				// get token
				String token = (String) invocation.getParameter(1);

				// The generic method adivce assumes to be called in the
				// afterRender render phase!
				// If it is called in a different render phase the XPath
				// expression may not work properly.
				MarkupWriter writer = (MarkupWriter) invocation.getParameter(0);
				try {
					List<Element> selectElements = TapestryXPath.xpath(
							".//form").selectElements(writer.getElement());
					for (Element formElement : selectElements) {
						addHiddenTokenToForm(formElement, token);
					}

					List<Element> linkElements = TapestryXPath.xpath(".//a")
							.selectElements(writer.getElement());

					for (Element linkElement : linkElements) {
						addParameterTokenToLink(linkElement, token);
					}

				} catch (JaxenException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
			}
		};
	}

	/**
	 * Returns a field handle for the given name. If it does not exist null is returned.
	 * @param plasticClass
	 * @param fieldName
	 * @return
	 */
	private FieldHandle getFieldHandleByName(PlasticClass plasticClass,
			String fieldName) {
		List<PlasticField> allFields = plasticClass.getAllFields();
		// Get a FieldHandle on the form Element
		for (PlasticField field : allFields) {
			if (field.getName().equals(fieldName)) {
				return field.getHandle();
			}
		}
		logger.debug("Field " + fieldName + " could not be found in class: "
				+ plasticClass.getClassName());
		return null;
	}

	/** 
	 * Checks if a given class is in the controlled packages of Tapestry.
	 * @param className
	 * @return
	 */
	private boolean isInControlledPackage(String className) {
		Collection<String> controlledPackages = componentClassResolver
				.getControlledPackageMapping().keySet();

		// Modify Tapestry core components
		for (String pack : controlledPackages) {
			if (className.startsWith(pack)) {
				return true;
			}
		}
		return false;
	}
}
